{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用 tf.data.Dataset 加载数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.2.0-rc4'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List 列表数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(1, shape=(), dtype=int32)\n",
      "tf.Tensor(2, shape=(), dtype=int32)\n",
      "tf.Tensor(3, shape=(), dtype=int32)\n"
     ]
    }
   ],
   "source": [
    "dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])\n",
    "for element in dataset:\n",
    "    print(element)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generator 生成器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "\n",
    "def gen():\n",
    "    for i in itertools.count(1):\n",
    "        yield (i, [1] * i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = tf.data.Dataset.from_generator(\n",
    "    gen,\n",
    "    (tf.int64, tf.int64),\n",
    "    (tf.TensorShape([]), tf.TensorShape([None])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(1, array([1])), (2, array([1, 1])), (3, array([1, 1, 1]))]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(dataset.take(3).as_numpy_iterator())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 文本文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "parent_dir = \"/work/chapter-2/files\"\n",
    "FILE_NAMES = ['cowper.txt', 'derby.txt', 'butler.txt']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "def labeler(example, index):\n",
    "    return example, tf.cast(index, tf.int64)  \n",
    "\n",
    "labeled_data_sets = []\n",
    "\n",
    "for i, file_name in enumerate(FILE_NAMES):\n",
    "    lines_dataset = tf.data.TextLineDataset(os.path.join(parent_dir, file_name))\n",
    "    labeled_dataset = lines_dataset.map(lambda ex: labeler(ex, i))\n",
    "    labeled_data_sets.append(labeled_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "BUFFER_SIZE = 50000\n",
    "BATCH_SIZE = 64\n",
    "TAKE_SIZE = 5000\n",
    "\n",
    "all_labeled_data = labeled_data_sets[0]\n",
    "for labeled_dataset in labeled_data_sets[1:]:\n",
    "    all_labeled_data = all_labeled_data.concatenate(labeled_dataset)\n",
    "\n",
    "all_labeled_data = all_labeled_data.shuffle(\n",
    "    BUFFER_SIZE, reshuffle_each_iteration=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(<tf.Tensor: shape=(), dtype=string, numpy=b\"While from the fight of men the Gods abstain'd,\">, <tf.Tensor: shape=(), dtype=int64, numpy=1>)\n",
      "(<tf.Tensor: shape=(), dtype=string, numpy=b'When thus ye shall have rallied every band'>, <tf.Tensor: shape=(), dtype=int64, numpy=0>)\n",
      "(<tf.Tensor: shape=(), dtype=string, numpy=b'My derivation? From the land I come'>, <tf.Tensor: shape=(), dtype=int64, numpy=0>)\n",
      "(<tf.Tensor: shape=(), dtype=string, numpy=b'and to attack us; therefore they shall be devoured of vultures; we'>, <tf.Tensor: shape=(), dtype=int64, numpy=2>)\n",
      "(<tf.Tensor: shape=(), dtype=string, numpy=b'While busied in such thought he stood, the ranks'>, <tf.Tensor: shape=(), dtype=int64, numpy=0>)\n"
     ]
    }
   ],
   "source": [
    "for ex in all_labeled_data.take(5):\n",
    "    print(ex)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
